from utils import *
from rtmlib import Body
from motionbert.mb_inferencer import motionbert_inferencer
from inference_2d_3d import inference_2d_3d_and_save
from tqdm import tqdm
import traceback
import time
import json

if __name__ == "__main__":
    date_string = time.strftime('%Y-%m-%d', time.localtime())
    log_path = f"log/inference_log_{date_string}.txt"
    logger = TextLogger(log_path)
    time_string = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    logger.log(f"start at: {time_string}")

    inferencer_2d = Body(
        pose='rtmo', # 使用 rtmpose 和 yolo detector 会报错
        to_openpose=False,
        mode='performance',  # balanced, performance, lightweight
        backend='onnxruntime',
        device='cuda')  # onnxruntime 最高支持cuda11.8
    
    config_path = os.path.join(get_project_root(), "motionbert/configs/pose3d/MB_train_h36m.yaml")
    ckpt_path = os.path.join(get_project_root(),
                             "motionbert/checkpoint/pose3d/MB_train_h36m_dual_loss/finish_at_2024_4_14/best_epoch.bin")
    inferencer_3d = motionbert_inferencer(config_path, ckpt_path)

    error_video_list = []

    video_dataset_root = os.path.join(get_project_root(), "datasets/emohugo_video")

    all_category = get_all_subfolders(video_dataset_root)
    all_videos = []
    for category in all_category:
        all_videos.extend(get_all_same_kind(category, [".avi"])) 

    # 断点继续信号
    continue_process = True

    for video in tqdm(all_videos, disable=False):
        # if video == "/home/hugoxana/Documents/emohugo/datasets/emohugo_video/anger/000018-S001-28.avi":
        #     continue_process = True

        if continue_process:
            print("processing ", video)
            try:
                inference_2d_3d_and_save(video, inferencer_2d, inferencer_3d, output_vis=True, silent=True)
            except KeyboardInterrupt:
                os._exit(0)                
            except:
                logger.log(f"ERROR in {video}")
                traceback.print_exc()
                with open(log_path, "a+") as f:
                    traceback.print_exc(file=f)
                error_video_list.append(video)
        else:
            print("skipping ", video)

    json_str = json.dumps(error_video_list, indent=4)
    error_video_path_file = f"log/error_video_path_{date_string}.json"
    with open(error_video_path_file,"w") as f:
        f.write(json_str)

    time_string = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
    logger.log(f"finish at: {time_string}")
    
